#!/usr/bin/python
# -*- coding: utf-8 -*-
# @Time: 20-8-19 下午10:47
# @Author: HuangChuQin
# @Site: 查询股票日数据并写入数据库
# @File: stock_data_append.py
# @Software: PyCharm

import time


class StockDataAppend:
    def __get_stock_info(self):

        rs = self.bs.query_stock_basic(self.code)

        print('query_stock_basic\n|-code:%s\n|-error_code:%s\n|-error_msg:%s' % (self.code, rs.error_code, rs.error_msg))

        return rs.data[0]

    def __get_stock_data(self, date=None):
        if self.is_init is False:
            # 只查当前时间
            date = time.strftime("%Y-%m-%d", time.localtime())
        elif date is None:
            print("初始化股票数据，必须传开始日期")
            return None
        rs = self.bs.query_history_k_data_plus(
            self.code, ','.join(list(self.dict.keys())), start_date=date, end_date=None, frequency="d"
        )
        print("query_history_k_data_plus\n|-error_code:%s\n|-error_msg:%s" % (rs.error_code, rs.error_msg))

        if len(rs.data) == 0:
            print("指定时间没有数据")
            return None

        data_list = []

        for datum in rs.data:
            data_dict = self.dict.copy()
            index = 0
            for key in data_dict:
                data_dict[key] = datum[index]
                index += 1
            data_list.append(data_dict)
        return data_list

    def __write_stock_data(self, file_name, data_list):
        folder = self.folder[file_name]
        folder.create_index([("date", 1)], unique=True)
        if self.is_init is True:
            folder.insert_many(data_list)
        else:
            folder.replace_one({"date": data_list[0]["date"]}, data_list[0], upsert=True)

    def __main(self):
        if self.date is not None:
            data_list = self.__get_stock_data(self.date)
        elif self.is_init is True:
            stock_info = self.__get_stock_info()
            data_list = self.__get_stock_data(stock_info[2])
        else:
            data_list = self.__get_stock_data(None)
        if data_list is None:
            return False
        self.__write_stock_data("stock_k_data_daily_" + self.code, data_list)
        return True

    def __init__(self, mg_folder, bao_stock, stock_code, field_dict, is_init=False, date=None):
        """
        查询股票日数据并写入数据库
        :param mg_folder: MongoDB 文档集合链接
        :param bao_stock: 已经登录处理的baostock链接
        :param stock_code: 股票代码
        :param field_dict: 股票指标字典
        :param is_init: 是否初始化股票仓库，初始化会爬取该股票的所有历史记录，默认为否，表示追加当天数据
        :param date: 指定追加的日期
        """
        self.folder = mg_folder
        self.bs = bao_stock
        self.code = stock_code
        self.dict = field_dict
        self.is_init = is_init
        self.date = date
        self.__main()
